from reid_libs.utils.using_utils.img_searchers import ImgSearcher
from reid_libs.datasets_manager import get_dataset
from reid_libs.utils.data_utils.person_transform import personTestTransform
from reid_libs.utils.data_utils.person_dataloader import PersonImageDataset
from torch.utils.data import DataLoader
from reid_libs.models import create_model
# from IPython import embed
import torch

# 使用阶段-选择设备
device = torch.device("cuda")
# device = torch.device("cpu")

# 使用阶段-超参数
BATCH_SIZE = 2
NUM_WORKERS = 0

# 使用的数据集
Test_Dataset = get_dataset('DebugDataset', r'D:\reid\DATA4REID')

query_loader = DataLoader(PersonImageDataset(Test_Dataset.query,
                                             transform=personTestTransform),
                          batch_size=BATCH_SIZE,
                          num_workers=NUM_WORKERS,
                          shuffle=False,
                          pin_memory=True)

gallery_loader = DataLoader(PersonImageDataset(Test_Dataset.gallery,
                                               transform=personTestTransform),
                            batch_size=BATCH_SIZE,
                            num_workers=NUM_WORKERS,
                            shuffle=False,
                            pin_memory=True)

# 使用的模型
model = create_model(name='resnet50_rga',
                     num_classes=751,
                     height=256,
                     width=128,
                     pretrained=True,
                     num_feat=2048)

# # 并行化
# model = torch.nn.DataParallel(model).to(device)

# 载入模型权重
resume = r'good_ckpt\market1501_resnet50rga_b64n4f2048\checkpoint_200.pth.tar'
print("Loading checkpoint from '{}'".format(resume))
# checkpoint = torch.load(resume, map_location='cpu')
checkpoint = torch.load(resume)
model.load_state_dict(checkpoint['state_dict'])
test_epoch = checkpoint['epoch']
print('test epoch:', test_epoch)

# 并行化
model = torch.nn.DataParallel(model).to(device)

# test/evaluate the model
searcher = ImgSearcher(model, selected_device=device)
# feat_ : 未归一化 ; feat : 归一化后
# feats_list = ['feat_', 'feat']
feats_list = ['feat']

searcher.use_worerank(
    query_loader,
    gallery_loader,
    Test_Dataset.query,
    Test_Dataset.gallery,
    #   metric=['euclidean', 'cosine'],
    metric=['cosine'],
    types_list=feats_list,
    top_k=20)
